# coding=utf-8
# Copyright 2024 The Uncertainty Baselines Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Utilities for multi-modal experiments."""
import jax
import jax.numpy as jnp


# Ported from
# https://github.com/google-research/big_vision/blob/main/big_vision/utils.py.
def bidirectional_contrastive_loss(zimg, ztxt, mask=None, reduction=False):
  """Bidirectional contrastive loss (e.g. for contrastive trainer/evaluator)."""
  # Note: in the original version of this function, there is a temperature
  # parameter `t` that is multiplied with the logits on the next line. This
  # temperature parameter is a model weight that is learnt via SGD. Here, we
  # assume that each zimg and ztxt are instead pre-multiplied with `sqrt(t)`.
  # This simplifies the code, removing the need to pass `t` around.
  logits = jnp.dot(zimg, ztxt.T)  # BF.FB = BB

  if mask is not None:
    # Set to negative infinity where mask = 0. Masked examples will disappear
    # under softmax (and argmax when calculating #correct).
    exclude = jnp.logical_not(mask)  # Now 1 if we don't want to keep.
    exclude = jnp.logical_or(exclude[:, None], exclude[None, :])
    logits = jnp.where(exclude, -jnp.inf, logits)

  # Note: assumed t is in a good range e.g. already passed through exp/softplus.
  l1 = -jnp.diag(jax.nn.log_softmax(logits, axis=1))  # NLL img->txt
  l2 = -jnp.diag(jax.nn.log_softmax(logits, axis=0))  # NLL txt->img
  l = 0.5 * (l1 + l2)

  if mask is not None:
    l = jnp.where(mask, l, 0)

  redux = jnp.mean if reduction else lambda x: x
  if reduction and mask is not None:
    redux = lambda x: jnp.sum(x * mask) / (jnp.sum(mask) + 1e-8)

  return redux(l), logits


# ImageNet-R and ImageNet-A label sets taken from
# https://github.com/google-research/robustness_metrics/blob/master/robustness_metrics/datasets/tfds.py.
_IMAGENET_R_LABELSET = [
    1, 2, 4, 6, 8, 9, 11, 13, 22, 23, 26, 29, 31, 39, 47, 63, 71, 76, 79, 84,
    90, 94, 96, 97, 99, 100, 105, 107, 113, 122, 125, 130, 132, 144, 145, 147,
    148, 150, 151, 155, 160, 161, 162, 163, 171, 172, 178, 187, 195, 199, 203,
    207, 208, 219, 231, 232, 234, 235, 242, 245, 247, 250, 251, 254, 259, 260,
    263, 265, 267, 269, 276, 277, 281, 288, 289, 291, 292, 293, 296, 299, 301,
    308, 309, 310, 311, 314, 315, 319, 323, 327, 330, 334, 335, 337, 338, 340,
    341, 344, 347, 353, 355, 361, 362, 365, 366, 367, 368, 372, 388, 390, 393,
    397, 401, 407, 413, 414, 425, 428, 430, 435, 437, 441, 447, 448, 457, 462,
    463, 469, 470, 471, 472, 476, 483, 487, 515, 546, 555, 558, 570, 579, 583,
    587, 593, 594, 596, 609, 613, 617, 621, 629, 637, 657, 658, 701, 717, 724,
    763, 768, 774, 776, 779, 780, 787, 805, 812, 815, 820, 824, 833, 847, 852,
    866, 875, 883, 889, 895, 907, 928, 931, 932, 933, 934, 936, 937, 943, 945,
    947, 948, 949, 951, 953, 954, 957, 963, 965, 967, 980, 981, 983, 988,
]

_IMAGENET_A_LABELSET = [
    6, 11, 13, 15, 17, 22, 23, 27, 30, 37, 39, 42, 47, 50, 57, 70, 71, 76, 79,
    89, 90, 94, 96, 97, 99, 105, 107, 108, 110, 113, 124, 125, 130, 132, 143,
    144, 150, 151, 207, 234, 235, 254, 277, 283, 287, 291, 295, 298, 301, 306,
    307, 308, 309, 310, 311, 313, 314, 315, 317, 319, 323, 324, 326, 327, 330,
    334, 335, 336, 347, 361, 363, 372, 378, 386, 397, 400, 401, 402, 404, 407,
    411, 416, 417, 420, 425, 428, 430, 437, 438, 445, 456, 457, 461, 462, 470,
    472, 483, 486, 488, 492, 496, 514, 516, 528, 530, 539, 542, 543, 549, 552,
    557, 561, 562, 569, 572, 573, 575, 579, 589, 606, 607, 609, 614, 626, 627,
    640, 641, 642, 643, 658, 668, 677, 682, 684, 687, 701, 704, 719, 736, 746,
    749, 752, 758, 763, 765, 768, 773, 774, 776, 779, 780, 786, 792, 797, 802,
    803, 804, 813, 815, 820, 823, 831, 833, 835, 839, 845, 847, 850, 859, 862,
    870, 879, 880, 888, 890, 897, 900, 907, 913, 924, 932, 933, 934, 937, 943,
    945, 947, 951, 954, 956, 957, 959, 971, 972, 980, 981, 984, 986, 987, 988,
]

# Class names and templates copied from:
# https://github.com/openai/CLIP/blob/main/notebooks/Prompt_Engineering_for_ImageNet.ipynb,
# https://github.com/openai/CLIP/blob/main/data/prompts.md, and
# https://arxiv.org/pdf/2111.07991.pdf (App I.2).
# Class names are only provided when the dataset is availble in TFDS.
# Note: some classname lists provided by OpenAI had an ordering which was not
# correct for TFDS and had to be reordered. Keep this in mind if modifying.
# TFDS ordering can be found with:
#   _, info = tfds.load(dataset_name, with_info=True)
#   [info.features['label'].int2str(i) for i in range(n_classes)]
_ZEROSHOT_CLASS_NAMES = {
    'imagenet': [
        'tench', 'goldfish', 'great white shark', 'tiger shark',
        'hammerhead shark', 'electric ray', 'stingray', 'rooster', 'hen',
        'ostrich', 'brambling', 'goldfinch', 'house finch', 'junco',
        'indigo bunting', 'American robin', 'bulbul', 'jay', 'magpie',
        'chickadee', 'American dipper', 'kite (bird of prey)', 'bald eagle',
        'vulture', 'great grey owl', 'fire salamander', 'smooth newt', 'newt',
        'spotted salamander', 'axolotl', 'American bullfrog', 'tree frog',
        'tailed frog', 'loggerhead sea turtle', 'leatherback sea turtle',
        'mud turtle', 'terrapin', 'box turtle', 'banded gecko', 'green iguana',
        'Carolina anole', 'desert grassland whiptail lizard', 'agama',
        'frilled-necked lizard', 'alligator lizard', 'Gila monster',
        'European green lizard', 'chameleon', 'Komodo dragon', 'Nile crocodile',
        'American alligator', 'triceratops', 'worm snake', 'ring-necked snake',
        'eastern hog-nosed snake', 'smooth green snake', 'kingsnake',
        'garter snake', 'water snake', 'vine snake', 'night snake',
        'boa constrictor', 'African rock python', 'Indian cobra', 'green mamba',
        'sea snake', 'Saharan horned viper', 'eastern diamondback rattlesnake',
        'sidewinder rattlesnake', 'trilobite', 'harvestman', 'scorpion',
        'yellow garden spider', 'barn spider', 'European garden spider',
        'southern black widow', 'tarantula', 'wolf spider', 'tick', 'centipede',
        'black grouse', 'ptarmigan', 'ruffed grouse', 'prairie grouse',
        'peafowl', 'quail', 'partridge', 'african grey parrot', 'macaw',
        'sulphur-crested cockatoo', 'lorikeet', 'coucal', 'bee eater',
        'hornbill', 'hummingbird', 'jacamar', 'toucan', 'duck',
        'red-breasted merganser', 'goose', 'black swan', 'tusker', 'echidna',
        'platypus', 'wallaby', 'koala', 'wombat', 'jellyfish', 'sea anemone',
        'brain coral', 'flatworm', 'nematode', 'conch', 'snail', 'slug',
        'sea slug', 'chiton', 'chambered nautilus', 'Dungeness crab',
        'rock crab', 'fiddler crab', 'red king crab', 'American lobster',
        'spiny lobster', 'crayfish', 'hermit crab', 'isopod', 'white stork',
        'black stork', 'spoonbill', 'flamingo', 'little blue heron',
        'great egret', 'bittern bird', 'crane bird', 'limpkin',
        'common gallinule', 'American coot', 'bustard', 'ruddy turnstone',
        'dunlin', 'common redshank', 'dowitcher', 'oystercatcher', 'pelican',
        'king penguin', 'albatross', 'grey whale', 'killer whale', 'dugong',
        'sea lion', 'Chihuahua', 'Japanese Chin', 'Maltese', 'Pekingese',
        'Shih Tzu', 'King Charles Spaniel', 'Papillon', 'toy terrier',
        'Rhodesian Ridgeback', 'Afghan Hound', 'Basset Hound', 'Beagle',
        'Bloodhound', 'Bluetick Coonhound', 'Black and Tan Coonhound',
        'Treeing Walker Coonhound', 'English foxhound', 'Redbone Coonhound',
        'borzoi', 'Irish Wolfhound', 'Italian Greyhound', 'Whippet',
        'Ibizan Hound', 'Norwegian Elkhound', 'Otterhound', 'Saluki',
        'Scottish Deerhound', 'Weimaraner', 'Staffordshire Bull Terrier',
        'American Staffordshire Terrier', 'Bedlington Terrier',
        'Border Terrier', 'Kerry Blue Terrier', 'Irish Terrier',
        'Norfolk Terrier', 'Norwich Terrier', 'Yorkshire Terrier',
        'Wire Fox Terrier', 'Lakeland Terrier', 'Sealyham Terrier',
        'Airedale Terrier', 'Cairn Terrier', 'Australian Terrier',
        'Dandie Dinmont Terrier', 'Boston Terrier', 'Miniature Schnauzer',
        'Giant Schnauzer', 'Standard Schnauzer', 'Scottish Terrier',
        'Tibetan Terrier', 'Australian Silky Terrier',
        'Soft-coated Wheaten Terrier', 'West Highland White Terrier',
        'Lhasa Apso', 'Flat-Coated Retriever', 'Curly-coated Retriever',
        'Golden Retriever', 'Labrador Retriever', 'Chesapeake Bay Retriever',
        'German Shorthaired Pointer', 'Vizsla', 'English Setter',
        'Irish Setter', 'Gordon Setter', 'Brittany dog', 'Clumber Spaniel',
        'English Springer Spaniel', 'Welsh Springer Spaniel', 'Cocker Spaniel',
        'Sussex Spaniel', 'Irish Water Spaniel', 'Kuvasz', 'Schipperke',
        'Groenendael dog', 'Malinois', 'Briard', 'Australian Kelpie',
        'Komondor', 'Old English Sheepdog', 'Shetland Sheepdog', 'collie',
        'Border Collie', 'Bouvier des Flandres dog', 'Rottweiler',
        'German Shepherd Dog', 'Dobermann', 'Miniature Pinscher',
        'Greater Swiss Mountain Dog', 'Bernese Mountain Dog',
        'Appenzeller Sennenhund', 'Entlebucher Sennenhund', 'Boxer',
        'Bullmastiff', 'Tibetan Mastiff', 'French Bulldog', 'Great Dane',
        'St. Bernard', 'husky', 'Alaskan Malamute', 'Siberian Husky',
        'Dalmatian', 'Affenpinscher', 'Basenji', 'pug', 'Leonberger',
        'Newfoundland dog', 'Great Pyrenees dog', 'Samoyed', 'Pomeranian',
        'Chow Chow', 'Keeshond', 'brussels griffon', 'Pembroke Welsh Corgi',
        'Cardigan Welsh Corgi', 'Toy Poodle', 'Miniature Poodle',
        'Standard Poodle', 'Mexican hairless dog (xoloitzcuintli)', 'grey wolf',
        'Alaskan tundra wolf', 'red wolf or maned wolf', 'coyote', 'dingo',
        'dhole', 'African wild dog', 'hyena', 'red fox', 'kit fox',
        'Arctic fox', 'grey fox', 'tabby cat', 'tiger cat', 'Persian cat',
        'Siamese cat', 'Egyptian Mau', 'cougar', 'lynx', 'leopard',
        'snow leopard', 'jaguar', 'lion', 'tiger', 'cheetah', 'brown bear',
        'American black bear', 'polar bear', 'sloth bear', 'mongoose',
        'meerkat', 'tiger beetle', 'ladybug', 'ground beetle',
        'longhorn beetle', 'leaf beetle', 'dung beetle', 'rhinoceros beetle',
        'weevil', 'fly', 'bee', 'ant', 'grasshopper', 'cricket insect',
        'stick insect', 'cockroach', 'praying mantis', 'cicada', 'leafhopper',
        'lacewing', 'dragonfly', 'damselfly', 'red admiral butterfly',
        'ringlet butterfly', 'monarch butterfly', 'small white butterfly',
        'sulphur butterfly', 'gossamer-winged butterfly', 'starfish',
        'sea urchin', 'sea cucumber', 'cottontail rabbit', 'hare',
        'Angora rabbit', 'hamster', 'porcupine', 'fox squirrel', 'marmot',
        'beaver', 'guinea pig', 'common sorrel horse', 'zebra', 'pig',
        'wild boar', 'warthog', 'hippopotamus', 'ox', 'water buffalo', 'bison',
        'ram (adult male sheep)', 'bighorn sheep', 'Alpine ibex', 'hartebeest',
        'impala (antelope)', 'gazelle', 'arabian camel', 'llama', 'weasel',
        'mink', 'European polecat', 'black-footed ferret', 'otter', 'skunk',
        'badger', 'armadillo', 'three-toed sloth', 'orangutan', 'gorilla',
        'chimpanzee', 'gibbon', 'siamang', 'guenon', 'patas monkey', 'baboon',
        'macaque', 'langur', 'black-and-white colobus', 'proboscis monkey',
        'marmoset', 'white-headed capuchin', 'howler monkey', 'titi monkey',
        "Geoffroy's spider monkey", 'common squirrel monkey',
        'ring-tailed lemur', 'indri', 'Asian elephant', 'African bush elephant',
        'red panda', 'giant panda', 'snoek fish', 'eel', 'silver salmon',
        'rock beauty fish', 'clownfish', 'sturgeon', 'gar fish', 'lionfish',
        'pufferfish', 'abacus', 'abaya', 'academic gown', 'accordion',
        'acoustic guitar', 'aircraft carrier', 'airliner', 'airship', 'altar',
        'ambulance', 'amphibious vehicle', 'analog clock', 'apiary', 'apron',
        'trash can', 'assault rifle', 'backpack', 'bakery', 'balance beam',
        'balloon', 'ballpoint pen', 'Band-Aid', 'banjo', 'baluster / handrail',
        'barbell', 'barber chair', 'barbershop', 'barn', 'barometer', 'barrel',
        'wheelbarrow', 'baseball', 'basketball', 'bassinet', 'bassoon',
        'swimming cap', 'bath towel', 'bathtub', 'station wagon', 'lighthouse',
        'beaker', 'military hat (bearskin or shako)', 'beer bottle',
        'beer glass', 'bell tower', 'baby bib', 'tandem bicycle', 'bikini',
        'ring binder', 'binoculars', 'birdhouse', 'boathouse', 'bobsleigh',
        'bolo tie', 'poke bonnet', 'bookcase', 'bookstore', 'bottle cap',
        'hunting bow', 'bow tie', 'brass memorial plaque', 'bra', 'breakwater',
        'breastplate', 'broom', 'bucket', 'buckle', 'bulletproof vest',
        'high-speed train', 'butcher shop', 'taxicab', 'cauldron', 'candle',
        'cannon', 'canoe', 'can opener', 'cardigan', 'car mirror', 'carousel',
        'tool kit', 'cardboard box / carton', 'car wheel',
        'automated teller machine', 'cassette', 'cassette player', 'castle',
        'catamaran', 'CD player', 'cello', 'mobile phone', 'chain',
        'chain-link fence', 'chain mail', 'chainsaw', 'storage chest',
        'chiffonier', 'bell or wind chime', 'china cabinet',
        'Christmas stocking', 'church', 'movie theater', 'cleaver',
        'cliff dwelling', 'cloak', 'clogs', 'cocktail shaker', 'coffee mug',
        'coffeemaker', 'spiral or coil', 'combination lock',
        'computer keyboard', 'candy store', 'container ship', 'convertible',
        'corkscrew', 'cornet', 'cowboy boot', 'cowboy hat', 'cradle',
        'construction crane', 'crash helmet', 'crate', 'infant bed',
        'Crock Pot', 'croquet ball', 'crutch', 'cuirass', 'dam', 'desk',
        'desktop computer', 'rotary dial telephone', 'diaper', 'digital clock',
        'digital watch', 'dining table', 'dishcloth', 'dishwasher',
        'disc brake', 'dock', 'dog sled', 'dome', 'doormat', 'drilling rig',
        'drum', 'drumstick', 'dumbbell', 'Dutch oven', 'electric fan',
        'electric guitar', 'electric locomotive', 'entertainment center',
        'envelope', 'espresso machine', 'face powder', 'feather boa',
        'filing cabinet', 'fireboat', 'fire truck', 'fire screen', 'flagpole',
        'flute', 'folding chair', 'football helmet', 'forklift', 'fountain',
        'fountain pen', 'four-poster bed', 'freight car', 'French horn',
        'frying pan', 'fur coat', 'garbage truck', 'gas mask or respirator',
        'gas pump', 'goblet', 'go-kart', 'golf ball', 'golf cart', 'gondola',
        'gong', 'gown', 'grand piano', 'greenhouse', 'radiator grille',
        'grocery store', 'guillotine', 'hair clip', 'hair spray', 'half-track',
        'hammer', 'hamper', 'hair dryer', 'hand-held computer', 'handkerchief',
        'hard disk drive', 'harmonica', 'harp', 'combine harvester', 'hatchet',
        'holster', 'home theater', 'honeycomb', 'hook', 'hoop skirt',
        'gymnastic horizontal bar', 'horse-drawn vehicle', 'hourglass', 'iPod',
        'clothes iron', 'carved pumpkin', 'jeans', 'jeep', 'T-shirt',
        'jigsaw puzzle', 'rickshaw', 'joystick', 'kimono', 'knee pad', 'knot',
        'lab coat', 'ladle', 'lampshade', 'laptop computer', 'lawn mower',
        'lens cap', 'letter opener', 'library', 'lifeboat', 'lighter',
        'limousine', 'ocean liner', 'lipstick', 'slip-on shoe', 'lotion',
        'music speaker', 'loupe magnifying glass', 'sawmill',
        'magnetic compass', 'messenger bag', 'mailbox', 'tights',
        'one-piece bathing suit', 'manhole cover', 'maraca', 'marimba', 'mask',
        'matchstick', 'maypole', 'maze', 'measuring cup', 'medicine cabinet',
        'megalith', 'microphone', 'microwave oven', 'military uniform',
        'milk can', 'minibus', 'miniskirt', 'minivan', 'missile', 'mitten',
        'mixing bowl', 'mobile home', 'ford model t', 'modem', 'monastery',
        'monitor', 'moped', 'mortar and pestle', 'graduation cap', 'mosque',
        'mosquito net', 'vespa', 'mountain bike', 'tent', 'computer mouse',
        'mousetrap', 'moving van', 'muzzle', 'metal nail', 'neck brace',
        'necklace', 'baby pacifier', 'notebook computer', 'obelisk', 'oboe',
        'ocarina', 'odometer', 'oil filter', 'pipe organ', 'oscilloscope',
        'overskirt', 'bullock cart', 'oxygen mask',
        'product packet / packaging', 'paddle', 'paddle wheel', 'padlock',
        'paintbrush', 'pajamas', 'palace', 'pan flute', 'paper towel',
        'parachute', 'parallel bars', 'park bench', 'parking meter',
        'railroad car', 'patio', 'payphone', 'pedestal', 'pencil case',
        'pencil sharpener', 'perfume', 'Petri dish', 'photocopier', 'plectrum',
        'Pickelhaube', 'picket fence', 'pickup truck', 'pier', 'piggy bank',
        'pill bottle', 'pillow', 'ping-pong ball', 'pinwheel', 'pirate ship',
        'drink pitcher', 'block plane', 'planetarium', 'plastic bag',
        'plate rack', 'farm plow', 'plunger', 'Polaroid camera', 'pole',
        'police van', 'poncho', 'pool table', 'soda bottle', 'plant pot',
        "potter's wheel", 'power drill', 'prayer rug', 'printer', 'prison',
        'missile', 'projector', 'hockey puck', 'punching bag', 'purse', 'quill',
        'quilt', 'race car', 'racket', 'radiator', 'radio', 'radio telescope',
        'rain barrel', 'recreational vehicle', 'fishing casting reel',
        'reflex camera', 'refrigerator', 'remote control', 'restaurant',
        'revolver', 'rifle', 'rocking chair', 'rotisserie', 'eraser',
        'rugby ball', 'ruler measuring stick', 'sneaker', 'safe', 'safety pin',
        'salt shaker', 'sandal', 'sarong', 'saxophone', 'scabbard',
        'weighing scale', 'school bus', 'schooner', 'scoreboard', 'CRT monitor',
        'screw', 'screwdriver', 'seat belt', 'sewing machine', 'shield',
        'shoe store', 'shoji screen / room divider', 'shopping basket',
        'shopping cart', 'shovel', 'shower cap', 'shower curtain', 'ski',
        'balaclava ski mask', 'sleeping bag', 'slide rule', 'sliding door',
        'slot machine', 'snorkel', 'snowmobile', 'snowplow', 'soap dispenser',
        'soccer ball', 'sock', 'solar thermal collector', 'sombrero',
        'soup bowl', 'keyboard space bar', 'space heater', 'space shuttle',
        'spatula', 'motorboat', 'spider web', 'spindle', 'sports car',
        'spotlight', 'stage', 'steam locomotive', 'through arch bridge',
        'steel drum', 'stethoscope', 'scarf', 'stone wall', 'stopwatch',
        'stove', 'strainer', 'tram', 'stretcher', 'couch', 'stupa', 'submarine',
        'suit', 'sundial', 'sunglasses', 'sunglasses', 'sunscreen',
        'suspension bridge', 'mop', 'sweatshirt', 'swim trunks / shorts',
        'swing', 'electrical switch', 'syringe', 'table lamp', 'tank',
        'tape player', 'teapot', 'teddy bear', 'television', 'tennis ball',
        'thatched roof', 'front curtain', 'thimble', 'threshing machine',
        'throne', 'tile roof', 'toaster', 'tobacco shop', 'toilet seat',
        'torch', 'totem pole', 'tow truck', 'toy store', 'tractor',
        'semi-trailer truck', 'tray', 'trench coat', 'tricycle', 'trimaran',
        'tripod', 'triumphal arch', 'trolleybus', 'trombone', 'hot tub',
        'turnstile', 'typewriter keyboard', 'umbrella', 'unicycle',
        'upright piano', 'vacuum cleaner', 'vase', 'vaulted or arched ceiling',
        'velvet fabric', 'vending machine', 'vestment', 'viaduct', 'violin',
        'volleyball', 'waffle iron', 'wall clock', 'wallet', 'wardrobe',
        'military aircraft', 'sink', 'washing machine', 'water bottle',
        'water jug', 'water tower', 'whiskey jug', 'whistle', 'hair wig',
        'window screen', 'window shade', 'Windsor tie', 'wine bottle',
        'airplane wing', 'wok', 'wooden spoon', 'wool', 'split-rail fence',
        'shipwreck', 'sailboat', 'yurt', 'website', 'comic book', 'crossword',
        'traffic or street sign', 'traffic light', 'dust jacket', 'menu',
        'plate', 'guacamole', 'consomme', 'hot pot', 'trifle', 'ice cream',
        'popsicle', 'baguette', 'bagel', 'pretzel', 'cheeseburger', 'hot dog',
        'mashed potatoes', 'cabbage', 'broccoli', 'cauliflower', 'zucchini',
        'spaghetti squash', 'acorn squash', 'butternut squash', 'cucumber',
        'artichoke', 'bell pepper', 'cardoon', 'mushroom', 'Granny Smith apple',
        'strawberry', 'orange', 'lemon', 'fig', 'pineapple', 'banana',
        'jackfruit', 'cherimoya (custard apple)', 'pomegranate', 'hay',
        'carbonara', 'chocolate syrup', 'dough', 'meatloaf', 'pizza', 'pot pie',
        'burrito', 'red wine', 'espresso', 'tea cup', 'eggnog', 'mountain',
        'bubble', 'cliff', 'coral reef', 'geyser', 'lakeshore', 'promontory',
        'sandbar', 'beach', 'valley', 'volcano', 'baseball player',
        'bridegroom', 'scuba diver', 'rapeseed', 'daisy',
        "yellow lady's slipper", 'corn', 'acorn', 'rose hip',
        'horse chestnut seed', 'coral fungus', 'agaric', 'gyromitra',
        'stinkhorn mushroom', 'earth star fungus', 'hen of the woods mushroom',
        'bolete', 'corn cob', 'toilet paper'
    ],
    'caltech101': [
        'accordion',
        'airplane',
        'anchor',
        'ant',
        'background',
        'barrel',
        'bass',
        'beaver',
        'binocular',
        'bonsai',
        'brain',
        'brontosaurus',
        'buddha',
        'butterfly',
        'camera',
        'cannon',
        'side of a car',
        'ceiling fan',
        'cellphone',
        'chair',
        'chandelier',
        'body of a cougar cat',
        'face of a cougar cat',
        'crab',
        'crayfish',
        'crocodile',
        'head of a  crocodile',
        'cup',
        'dalmatian',
        'dollar bill',
        'dolphin',
        'dragonfly',
        'electric guitar',
        'elephant',
        'emu',
        'euphonium',
        'ewer',
        'off-center face',
        'centered face',
        'ferry',
        'flamingo',
        'head of a flamingo',
        'garfield',
        'gerenuk',
        'gramophone',
        'grand piano',
        'hawksbill',
        'headphone',
        'hedgehog',
        'helicopter',
        'ibis',
        'inline skate',
        'joshua tree',
        'kangaroo',
        'ketch',
        'lamp',
        'laptop',
        'leopard',
        'llama',
        'lobster',
        'lotus',
        'mandolin',
        'mayfly',
        'menorah',
        'metronome',
        'minaret',
        'motorbike',
        'nautilus',
        'octopus',
        'okapi',
        'pagoda',
        'panda',
        'pigeon',
        'pizza',
        'platypus',
        'pyramid',
        'revolver',
        'rhino',
        'rooster',
        'saxophone',
        'schooner',
        'scissors',
        'scorpion',
        'sea horse',
        'snoopy (cartoon beagle)',
        'soccer ball',
        'stapler',
        'starfish',
        'stegosaurus',
        'stop sign',
        'strawberry',
        'sunflower',
        'tick',
        'trilobite',
        'umbrella',
        'watch',
        'water lilly',
        'wheelchair',
        'wild cat',
        'windsor chair',
        'wrench',
        'yin and yang symbol',
    ],
    'cars196': [
        'AM General Hummer SUV 2000',
        'Acura RL Sedan 2012',
        'Acura TL Sedan 2012',
        'Acura TL Type-S 2008',
        'Acura TSX Sedan 2012',
        'Acura Integra Type R 2001',
        'Acura ZDX Hatchback 2012',
        'Aston Martin V8 Vantage Convertible 2012',
        'Aston Martin V8 Vantage Coupe 2012',
        'Aston Martin Virage Convertible 2012',
        'Aston Martin Virage Coupe 2012',
        'Audi RS 4 Convertible 2008',
        'Audi A5 Coupe 2012',
        'Audi TTS Coupe 2012',
        'Audi R8 Coupe 2012',
        'Audi V8 Sedan 1994',
        'Audi 100 Sedan 1994',
        'Audi 100 Wagon 1994',
        'Audi TT Hatchback 2011',
        'Audi S6 Sedan 2011',
        'Audi S5 Convertible 2012',
        'Audi S5 Coupe 2012',
        'Audi S4 Sedan 2012',
        'Audi S4 Sedan 2007',
        'Audi TT RS Coupe 2012',
        'BMW ActiveHybrid 5 Sedan 2012',
        'BMW 1 Series Convertible 2012',
        'BMW 1 Series Coupe 2012',
        'BMW 3 Series Sedan 2012',
        'BMW 3 Series Wagon 2012',
        'BMW 6 Series Convertible 2007',
        'BMW X5 SUV 2007',
        'BMW X6 SUV 2012',
        'BMW M3 Coupe 2012',
        'BMW M5 Sedan 2010',
        'BMW M6 Convertible 2010',
        'BMW X3 SUV 2012',
        'BMW Z4 Convertible 2012',
        'Bentley Continental Supersports Conv. Convertible 2012',
        'Bentley Arnage Sedan 2009',
        'Bentley Mulsanne Sedan 2011',
        'Bentley Continental GT Coupe 2012',
        'Bentley Continental GT Coupe 2007',
        'Bentley Continental Flying Spur Sedan 2007',
        'Bugatti Veyron 16.4 Convertible 2009',
        'Bugatti Veyron 16.4 Coupe 2009',
        'Buick Regal GS 2012',
        'Buick Rainier SUV 2007',
        'Buick Verano Sedan 2012',
        'Buick Enclave SUV 2012',
        'Cadillac CTS-V Sedan 2012',
        'Cadillac SRX SUV 2012',
        'Cadillac Escalade EXT Crew Cab 2007',
        'Chevrolet Silverado 1500 Hybrid Crew Cab 2012',
        'Chevrolet Corvette Convertible 2012',
        'Chevrolet Corvette ZR1 2012',
        'Chevrolet Corvette Ron Fellows Edition Z06 2007',
        'Chevrolet Traverse SUV 2012',
        'Chevrolet Camaro Convertible 2012',
        'Chevrolet HHR SS 2010',
        'Chevrolet Impala Sedan 2007',
        'Chevrolet Tahoe Hybrid SUV 2012',
        'Chevrolet Sonic Sedan 2012',
        'Chevrolet Express Cargo Van 2007',
        'Chevrolet Avalanche Crew Cab 2012',
        'Chevrolet Cobalt SS 2010',
        'Chevrolet Malibu Hybrid Sedan 2010',
        'Chevrolet TrailBlazer SS 2009',
        'Chevrolet Silverado 2500HD Regular Cab 2012',
        'Chevrolet Silverado 1500 Classic Extended Cab 2007',
        'Chevrolet Express Van 2007',
        'Chevrolet Monte Carlo Coupe 2007',
        'Chevrolet Malibu Sedan 2007',
        'Chevrolet Silverado 1500 Extended Cab 2012',
        'Chevrolet Silverado 1500 Regular Cab 2012',
        'Chrysler Aspen SUV 2009',
        'Chrysler Sebring Convertible 2010',
        'Chrysler Town and Country Minivan 2012',
        'Chrysler 300 SRT-8 2010',
        'Chrysler Crossfire Convertible 2008',
        'Chrysler PT Cruiser Convertible 2008',
        'Daewoo Nubira Wagon 2002',
        'Dodge Caliber Wagon 2012',
        'Dodge Caliber Wagon 2007',
        'Dodge Caravan Minivan 1997',
        'Dodge Ram Pickup 3500 Crew Cab 2010',
        'Dodge Ram Pickup 3500 Quad Cab 2009',
        'Dodge Sprinter Cargo Van 2009',
        'Dodge Journey SUV 2012',
        'Dodge Dakota Crew Cab 2010',
        'Dodge Dakota Club Cab 2007',
        'Dodge Magnum Wagon 2008',
        'Dodge Challenger SRT8 2011',
        'Dodge Durango SUV 2012',
        'Dodge Durango SUV 2007',
        'Dodge Charger Sedan 2012',
        'Dodge Charger SRT-8 2009',
        'Eagle Talon Hatchback 1998',
        'FIAT 500 Abarth 2012',
        'FIAT 500 Convertible 2012',
        'Ferrari FF Coupe 2012',
        'Ferrari California Convertible 2012',
        'Ferrari 458 Italia Convertible 2012',
        'Ferrari 458 Italia Coupe 2012',
        'Fisker Karma Sedan 2012',
        'Ford F-450 Super Duty Crew Cab 2012',
        'Ford Mustang Convertible 2007',
        'Ford Freestar Minivan 2007',
        'Ford Expedition EL SUV 2009',
        'Ford Edge SUV 2012',
        'Ford Ranger SuperCab 2011',
        'Ford GT Coupe 2006',
        'Ford F-150 Regular Cab 2012',
        'Ford F-150 Regular Cab 2007',
        'Ford Focus Sedan 2007',
        'Ford E-Series Wagon Van 2012',
        'Ford Fiesta Sedan 2012',
        'GMC Terrain SUV 2012',
        'GMC Savana Van 2012',
        'GMC Yukon Hybrid SUV 2012',
        'GMC Acadia SUV 2012',
        'GMC Canyon Extended Cab 2012',
        'Geo Metro Convertible 1993',
        'HUMMER H3T Crew Cab 2010',
        'HUMMER H2 SUT Crew Cab 2009',
        'Honda Odyssey Minivan 2012',
        'Honda Odyssey Minivan 2007',
        'Honda Accord Coupe 2012',
        'Honda Accord Sedan 2012',
        'Hyundai Veloster Hatchback 2012',
        'Hyundai Santa Fe SUV 2012',
        'Hyundai Tucson SUV 2012',
        'Hyundai Veracruz SUV 2012',
        'Hyundai Sonata Hybrid Sedan 2012',
        'Hyundai Elantra Sedan 2007',
        'Hyundai Accent Sedan 2012',
        'Hyundai Genesis Sedan 2012',
        'Hyundai Sonata Sedan 2012',
        'Hyundai Elantra Touring Hatchback 2012',
        'Hyundai Azera Sedan 2012',
        'Infiniti G Coupe IPL 2012',
        'Infiniti QX56 SUV 2011',
        'Isuzu Ascender SUV 2008',
        'Jaguar XK XKR 2012',
        'Jeep Patriot SUV 2012',
        'Jeep Wrangler SUV 2012',
        'Jeep Liberty SUV 2012',
        'Jeep Grand Cherokee SUV 2012',
        'Jeep Compass SUV 2012',
        'Lamborghini Reventon Coupe 2008',
        'Lamborghini Aventador Coupe 2012',
        'Lamborghini Gallardo LP 570-4 Superleggera 2012',
        'Lamborghini Diablo Coupe 2001',
        'Land Rover Range Rover SUV 2012',
        'Land Rover LR2 SUV 2012',
        'Lincoln Town Car Sedan 2011',
        'MINI Cooper Roadster Convertible 2012',
        'Maybach Landaulet Convertible 2012',
        'Mazda Tribute SUV 2011',
        'McLaren MP4-12C Coupe 2012',
        'Mercedes-Benz 300-Class Convertible 1993',
        'Mercedes-Benz C-Class Sedan 2012',
        'Mercedes-Benz SL-Class Coupe 2009',
        'Mercedes-Benz E-Class Sedan 2012',
        'Mercedes-Benz S-Class Sedan 2012',
        'Mercedes-Benz Sprinter Van 2012',
        'Mitsubishi Lancer Sedan 2012',
        'Nissan Leaf Hatchback 2012',
        'Nissan NV Passenger Van 2012',
        'Nissan Juke Hatchback 2012',
        'Nissan 240SX Coupe 1998',
        'Plymouth Neon Coupe 1999',
        'Porsche Panamera Sedan 2012',
        'Ram C/V Cargo Van Minivan 2012',
        'Rolls-Royce Phantom Drophead Coupe Convertible 2012',
        'Rolls-Royce Ghost Sedan 2012',
        'Rolls-Royce Phantom Sedan 2012',
        'Scion xD Hatchback 2012',
        'Spyker C8 Convertible 2009',
        'Spyker C8 Coupe 2009',
        'Suzuki Aerio Sedan 2007',
        'Suzuki Kizashi Sedan 2012',
        'Suzuki SX4 Hatchback 2012',
        'Suzuki SX4 Sedan 2012',
        'Tesla Model S Sedan 2012',
        'Toyota Sequoia SUV 2012',
        'Toyota Camry Sedan 2012',
        'Toyota Corolla Sedan 2012',
        'Toyota 4Runner SUV 2012',
        'Volkswagen Golf Hatchback 2012',
        'Volkswagen Golf Hatchback 1991',
        'Volkswagen Beetle Hatchback 2012',
        'Volvo C30 Hatchback 2012',
        'Volvo 240 Sedan 1993',
        'Volvo XC90 SUV 2007',
        'smart fortwo Convertible 2012',
    ],
    'cifar10': [
        'airplane',
        'automobile',
        'bird',
        'cat',
        'deer',
        'dog',
        'frog',
        'horse',
        'ship',
        'truck',
    ],
    'cifar100': [
        'apple',
        'aquarium fish',
        'baby',
        'bear',
        'beaver',
        'bed',
        'bee',
        'beetle',
        'bicycle',
        'bottle',
        'bowl',
        'boy',
        'bridge',
        'bus',
        'butterfly',
        'camel',
        'can',
        'castle',
        'caterpillar',
        'cattle',
        'chair',
        'chimpanzee',
        'clock',
        'cloud',
        'cockroach',
        'couch',
        'crab',
        'crocodile',
        'cup',
        'dinosaur',
        'dolphin',
        'elephant',
        'flatfish',
        'forest',
        'fox',
        'girl',
        'hamster',
        'house',
        'kangaroo',
        'keyboard',
        'lamp',
        'lawn mower',
        'leopard',
        'lion',
        'lizard',
        'lobster',
        'man',
        'maple tree',
        'motorcycle',
        'mountain',
        'mouse',
        'mushroom',
        'oak tree',
        'orange',
        'orchid',
        'otter',
        'palm tree',
        'pear',
        'pickup truck',
        'pine tree',
        'plain',
        'plate',
        'poppy',
        'porcupine',
        'possum',
        'rabbit',
        'raccoon',
        'ray',
        'road',
        'rocket',
        'rose',
        'sea',
        'seal',
        'shark',
        'shrew',
        'skunk',
        'skyscraper',
        'snail',
        'snake',
        'spider',
        'squirrel',
        'streetcar',
        'sunflower',
        'sweet pepper',
        'table',
        'tank',
        'telephone',
        'television',
        'tiger',
        'tractor',
        'train',
        'trout',
        'tulip',
        'turtle',
        'wardrobe',
        'whale',
        'willow tree',
        'wolf',
        'woman',
        'worm',
    ],
    'dtd': [
        'banded',
        'blotchy',
        'braided',
        'bubbly',
        'bumpy',
        'chequered',
        'cobwebbed',
        'cracked',
        'crosshatched',
        'crystalline',
        'dotted',
        'fibrous',
        'flecked',
        'freckled',
        'frilly',
        'gauzy',
        'grid',
        'grooved',
        'honeycombed',
        'interlaced',
        'knitted',
        'lacelike',
        'lined',
        'marbled',
        'matted',
        'meshed',
        'paisley',
        'perforated',
        'pitted',
        'pleated',
        'polka-dotted',
        'porous',
        'potholed',
        'scaly',
        'smeared',
        'spiralled',
        'sprinkled',
        'stained',
        'stratified',
        'striped',
        'studded',
        'swirly',
        'veined',
        'waffled',
        'woven',
        'wrinkled',
        'zigzagged',
    ],
    'eurosat': [
        'annual crop land',
        'forest',
        'brushland or shrubland',
        'highway or road',
        'industrial buildings or commercial buildings',
        'pasture land',
        'permanent crop land',
        'residential buildings or homes or apartments',
        'river',
        'lake or sea',
    ],
    'food101': [
        'apple pie',
        'baby back ribs',
        'baklava',
        'beef carpaccio',
        'beef tartare',
        'beet salad',
        'beignets',
        'bibimbap',
        'bread pudding',
        'breakfast burrito',
        'bruschetta',
        'caesar salad',
        'cannoli',
        'caprese salad',
        'carrot cake',
        'ceviche',
        'cheesecake',
        'cheese plate',
        'chicken curry',
        'chicken quesadilla',
        'chicken wings',
        'chocolate cake',
        'chocolate mousse',
        'churros',
        'clam chowder',
        'club sandwich',
        'crab cakes',
        'creme brulee',
        'croque madame',
        'cup cakes',
        'deviled eggs',
        'donuts',
        'dumplings',
        'edamame',
        'eggs benedict',
        'escargots',
        'falafel',
        'filet mignon',
        'fish and chips',
        'foie gras',
        'french fries',
        'french onion soup',
        'french toast',
        'fried calamari',
        'fried rice',
        'frozen yogurt',
        'garlic bread',
        'gnocchi',
        'greek salad',
        'grilled cheese sandwich',
        'grilled salmon',
        'guacamole',
        'gyoza',
        'hamburger',
        'hot and sour soup',
        'hot dog',
        'huevos rancheros',
        'hummus',
        'ice cream',
        'lasagna',
        'lobster bisque',
        'lobster roll sandwich',
        'macaroni and cheese',
        'macarons',
        'miso soup',
        'mussels',
        'nachos',
        'omelette',
        'onion rings',
        'oysters',
        'pad thai',
        'paella',
        'pancakes',
        'panna cotta',
        'peking duck',
        'pho',
        'pizza',
        'pork chop',
        'poutine',
        'prime rib',
        'pulled pork sandwich',
        'ramen',
        'ravioli',
        'red velvet cake',
        'risotto',
        'samosa',
        'sashimi',
        'scallops',
        'seaweed salad',
        'shrimp and grits',
        'spaghetti bolognese',
        'spaghetti carbonara',
        'spring rolls',
        'steak',
        'strawberry shortcake',
        'sushi',
        'tacos',
        'takoyaki',
        'tiramisu',
        'tuna tartare',
        'waffles',
    ],
    'mnist': [
        '0',
        '1',
        '2',
        '3',
        '4',
        '5',
        '6',
        '7',
        '8',
        '9',
    ],
    'oxford_flowers102': [
        'pink primrose',
        'hard-leaved pocket orchid',
        'canterbury bells',
        'sweet pea',
        'english marigold',
        'tiger lily',
        'moon orchid',
        'bird of paradise',
        'monkshood',
        'globe thistle',
        'snapdragon',
        "colt's foot",
        'king protea',
        'spear thistle',
        'yellow iris',
        'globe flower',
        'purple coneflower',
        'peruvian lily',
        'balloon flower',
        'giant white arum lily',
        'fire lily',
        'pincushion flower',
        'fritillary',
        'red ginger',
        'grape hyacinth',
        'corn poppy',
        'prince of wales feathers',
        'stemless gentian',
        'artichoke',
        'sweet william',
        'carnation',
        'garden phlox',
        'love in the mist',
        'mexican aster',
        'alpine sea holly',
        'ruby-lipped cattleya',
        'cape flower',
        'great masterwort',
        'siam tulip',
        'lenten rose',
        'barbeton daisy',
        'daffodil',
        'sword lily',
        'poinsettia',
        'bolero deep blue',
        'wallflower',
        'marigold',
        'buttercup',
        'oxeye daisy',
        'common dandelion',
        'petunia',
        'wild pansy',
        'primula',
        'sunflower',
        'pelargonium',
        'bishop of llandaff',
        'gaura',
        'geranium',
        'orange dahlia',
        'pink and yellow dahlia',
        'cautleya spicata',
        'japanese anemone',
        'black-eyed susan',
        'silverbush',
        'californian poppy',
        'osteospermum',
        'spring crocus',
        'bearded iris',
        'windflower',
        'tree poppy',
        'gazania',
        'azalea',
        'water lily',
        'rose',
        'thorn apple',
        'morning glory',
        'passion flower',
        'lotus',
        'toad lily',
        'anthurium',
        'frangipani',
        'clematis',
        'hibiscus',
        'columbine',
        'desert-rose',
        'tree mallow',
        'magnolia',
        'cyclamen',
        'watercress',
        'canna lily',
        'hippeastrum',
        'bee balm',
        'air plant',
        'foxglove',
        'bougainvillea',
        'camellia',
        'mallow',
        'mexican petunia',
        'bromelia',
        'blanket flower',
        'trumpet creeper',
        'blackberry lily',
    ],
    'oxford_iiit_pet': [
        'Abyssinian',
        'american bulldog',
        'american pit bull terrier',
        'basset hound',
        'beagle',
        'Bengal',
        'Birman',
        'Bombay',
        'boxer',
        'British Shorthair',
        'chihuahua',
        'Egyptian Mau',
        'english cocker spaniel',
        'english setter',
        'german shorthaired',
        'great pyrenees',
        'havanese',
        'japanese chin',
        'keeshond',
        'leonberger',
        'Maine Coon',
        'miniature pinscher',
        'newfoundland',
        'Persian',
        'pomeranian',
        'pug',
        'Ragdoll',
        'Russian Blue',
        'saint bernard',
        'samoyed',
        'scottish terrier',
        'shiba inu',
        'Siamese',
        'Sphynx',
        'staffordshire bull terrier',
        'wheaten terrier',
        'yorkshire terrier',
    ],
    'patch_camelyon': [
        'lymph node',
        'lymph node containing metastatic tumor tissue',
    ],
    'resisc45': [
        'airplane',
        'airport',
        'baseball diamond',
        'basketball court',
        'beach',
        'bridge',
        'chaparral',
        'church',
        'circular farmland',
        'cloud',
        'commercial area',
        'dense residential',
        'desert',
        'forest',
        'freeway',
        'golf course',
        'ground track field',
        'harbor',
        'industrial area',
        'intersection',
        'island',
        'lake',
        'meadow',
        'medium residential',
        'mobile home park',
        'mountain',
        'overpass',
        'palace',
        'parking lot',
        'railway',
        'railway station',
        'rectangular farmland',
        'river',
        'roundabout',
        'runway',
        'sea ice',
        'ship',
        'snowberg',
        'sparse residential',
        'stadium',
        'storage tank',
        'tennis court',
        'terrace',
        'thermal power station',
        'wetland',
    ],
    'stl10': [
        'airplane',
        'bird',
        'car',
        'cat',
        'deer',
        'dog',
        'horse',
        'monkey',
        'ship',
        'truck',
    ],
    'sun397': [
        'abbey',
        'airplane cabin',
        'airport terminal',
        'alley',
        'amphitheater',
        'amusement arcade',
        'amusement park',
        'anechoic chamber',
        'apartment building outdoor',
        'apse indoor',
        'aquarium',
        'aqueduct',
        'arch',
        'archive',
        'arrival gate outdoor',
        'art gallery',
        'art school',
        'art studio',
        'assembly line',
        'athletic field outdoor',
        'atrium public',
        'attic',
        'auditorium',
        'auto factory',
        'badlands',
        'badminton court indoor',
        'baggage claim',
        'bakery shop',
        'balcony exterior',
        'balcony interior',
        'ball pit',
        'ballroom',
        'bamboo forest',
        'banquet hall',
        'bar',
        'barn',
        'barndoor',
        'baseball field',
        'basement',
        'basilica',
        'basketball court outdoor',
        'bathroom',
        'batters box',
        'bayou',
        'bazaar indoor',
        'bazaar outdoor',
        'beach',
        'beauty salon',
        'bedroom',
        'berth',
        'biology laboratory',
        'bistro indoor',
        'boardwalk',
        'boat deck',
        'boathouse',
        'bookstore',
        'booth indoor',
        'botanical garden',
        'bow window indoor',
        'bow window outdoor',
        'bowling alley',
        'boxing ring',
        'brewery indoor',
        'bridge',
        'building facade',
        'bullring',
        'burial chamber',
        'bus interior',
        'butchers shop',
        'butte',
        'cabin outdoor',
        'cafeteria',
        'campsite',
        'campus',
        'canal natural',
        'canal urban',
        'candy store',
        'canyon',
        'car interior backseat',
        'car interior frontseat',
        'carrousel',
        'casino indoor',
        'castle',
        'catacomb',
        'cathedral indoor',
        'cathedral outdoor',
        'cavern indoor',
        'cemetery',
        'chalet',
        'cheese factory',
        'chemistry lab',
        'chicken coop indoor',
        'chicken coop outdoor',
        'childs room',
        'church indoor',
        'church outdoor',
        'classroom',
        'clean room',
        'cliff',
        'cloister indoor',
        'closet',
        'clothing store',
        'coast',
        'cockpit',
        'coffee shop',
        'computer room',
        'conference center',
        'conference room',
        'construction site',
        'control room',
        'control tower outdoor',
        'corn field',
        'corral',
        'corridor',
        'cottage garden',
        'courthouse',
        'courtroom',
        'courtyard',
        'covered bridge exterior',
        'creek',
        'crevasse',
        'crosswalk',
        'cubicle office',
        'dam',
        'delicatessen',
        'dentists office',
        'desert sand',
        'desert vegetation',
        'diner indoor',
        'diner outdoor',
        'dinette home',
        'dinette vehicle',
        'dining car',
        'dining room',
        'discotheque',
        'dock',
        'doorway outdoor',
        'dorm room',
        'driveway',
        'driving range outdoor',
        'drugstore',
        'electrical substation',
        'elevator door',
        'elevator interior',
        'elevator shaft',
        'engine room',
        'escalator indoor',
        'excavation',
        'factory indoor',
        'fairway',
        'fastfood restaurant',
        'field cultivated',
        'field wild',
        'fire escape',
        'fire station',
        'firing range indoor',
        'fishpond',
        'florist shop indoor',
        'food court',
        'forest broadleaf',
        'forest needleleaf',
        'forest path',
        'forest road',
        'formal garden',
        'fountain',
        'galley',
        'game room',
        'garage indoor',
        'garbage dump',
        'gas station',
        'gazebo exterior',
        'general store indoor',
        'general store outdoor',
        'gift shop',
        'golf course',
        'greenhouse indoor',
        'greenhouse outdoor',
        'gymnasium indoor',
        'hangar indoor',
        'hangar outdoor',
        'harbor',
        'hayfield',
        'heliport',
        'herb garden',
        'highway',
        'hill',
        'home office',
        'hospital',
        'hospital room',
        'hot spring',
        'hot tub outdoor',
        'hotel outdoor',
        'hotel room',
        'house',
        'hunting lodge outdoor',
        'ice cream parlor',
        'ice floe',
        'ice shelf',
        'ice skating rink indoor',
        'ice skating rink outdoor',
        'iceberg',
        'igloo',
        'industrial area',
        'inn outdoor',
        'islet',
        'jacuzzi indoor',
        'jail cell',
        'jail indoor',
        'jewelry shop',
        'kasbah',
        'kennel indoor',
        'kennel outdoor',
        'kindergarden classroom',
        'kitchen',
        'kitchenette',
        'labyrinth outdoor',
        'lake natural',
        'landfill',
        'landing deck',
        'laundromat',
        'lecture room',
        'library indoor',
        'library outdoor',
        'lido deck outdoor',
        'lift bridge',
        'lighthouse',
        'limousine interior',
        'living room',
        'lobby',
        'lock chamber',
        'locker room',
        'mansion',
        'manufactured home',
        'market indoor',
        'market outdoor',
        'marsh',
        'martial arts gym',
        'mausoleum',
        'medina',
        'moat water',
        'monastery outdoor',
        'mosque indoor',
        'mosque outdoor',
        'motel',
        'mountain',
        'mountain snowy',
        'movie theater indoor',
        'museum indoor',
        'music store',
        'music studio',
        'nuclear power plant outdoor',
        'nursery',
        'oast house',
        'observatory outdoor',
        'ocean',
        'office',
        'office building',
        'oil refinery outdoor',
        'oilrig',
        'operating room',
        'orchard',
        'outhouse outdoor',
        'pagoda',
        'palace',
        'pantry',
        'park',
        'parking garage indoor',
        'parking garage outdoor',
        'parking lot',
        'parlor',
        'pasture',
        'patio',
        'pavilion',
        'pharmacy',
        'phone booth',
        'physics laboratory',
        'picnic area',
        'pilothouse indoor',
        'planetarium outdoor',
        'playground',
        'playroom',
        'plaza',
        'podium indoor',
        'podium outdoor',
        'pond',
        'poolroom establishment',
        'poolroom home',
        'power plant outdoor',
        'promenade deck',
        'pub indoor',
        'pulpit',
        'putting green',
        'racecourse',
        'raceway',
        'raft',
        'railroad track',
        'rainforest',
        'reception',
        'recreation room',
        'residential neighborhood',
        'restaurant',
        'restaurant kitchen',
        'restaurant patio',
        'rice paddy',
        'riding arena',
        'river',
        'rock arch',
        'rope bridge',
        'ruin',
        'runway',
        'sandbar',
        'sandbox',
        'sauna',
        'schoolhouse',
        'sea cliff',
        'server room',
        'shed',
        'shoe shop',
        'shopfront',
        'shopping mall indoor',
        'shower',
        'skatepark',
        'ski lodge',
        'ski resort',
        'ski slope',
        'sky',
        'skyscraper',
        'slum',
        'snowfield',
        'squash court',
        'stable',
        'stadium baseball',
        'stadium football',
        'stage indoor',
        'staircase',
        'street',
        'subway interior',
        'subway station platform',
        'supermarket',
        'sushi bar',
        'swamp',
        'swimming pool indoor',
        'swimming pool outdoor',
        'synagogue indoor',
        'synagogue outdoor',
        'television studio',
        'temple east asia',
        'temple south asia',
        'tennis court indoor',
        'tennis court outdoor',
        'tent outdoor',
        'theater indoor procenium',
        'theater indoor seats',
        'thriftshop',
        'throne room',
        'ticket booth',
        'toll plaza',
        'topiary garden',
        'tower',
        'toyshop',
        'track outdoor',
        'train railway',
        'train station platform',
        'tree farm',
        'tree house',
        'trench',
        'underwater coral reef',
        'utility room',
        'valley',
        'van interior',
        'vegetable garden',
        'veranda',
        'veterinarians office',
        'viaduct',
        'videostore',
        'village',
        'vineyard',
        'volcano',
        'volleyball court indoor',
        'volleyball court outdoor',
        'waiting room',
        'warehouse indoor',
        'water tower',
        'waterfall block',
        'waterfall fan',
        'waterfall plunge',
        'watering hole',
        'wave',
        'wet bar',
        'wheat field',
        'wind farm',
        'windmill',
        'wine cellar barrel storage',
        'wine cellar bottle storage',
        'wrestling ring indoor',
        'yard',
        'youth hostel',
    ],
    'svhn': [
        'zero',
        'one',
        'two',
        'three',
        'four',
        'five',
        'six',
        'seven',
        'eight',
        'nine',
    ],
    'ucf101': [
        'Apply Eye Makeup',
        'Apply Lipstick',
        'Archery',
        'Baby Crawling',
        'Balance Beam',
        'Band Marching',
        'Baseball Pitch',
        'Basketball',
        'Basketball Dunk',
        'Bench Press',
        'Biking',
        'Billiards',
        'Blow Dry Hair',
        'Blowing Candles',
        'Body Weight Squats',
        'Bowling',
        'Boxing Punching Bag',
        'Boxing Speed Bag',
        'Breast Stroke',
        'Brushing Teeth',
        'Clean And Jerk',
        'Cliff Diving',
        'Cricket Bowling',
        'Cricket Shot',
        'Cutting In Kitchen',
        'Diving',
        'Drumming',
        'Fencing',
        'Field Hockey Penalty',
        'Floor Gymnastics',
        'Frisbee Catch',
        'Front Crawl',
        'Golf Swing',
        'Haircut',
        'Hammering',
        'Hammer Throw',
        'Hand Stand Pushups',
        'Handstand Walking',
        'Head Massage',
        'High Jump',
        'Horse Race',
        'Horse Riding',
        'Hula Hoop',
        'Ice Dancing',
        'Javelin Throw',
        'Juggling Balls',
        'Jumping Jack',
        'Jump Rope',
        'Kayaking',
        'Knitting',
        'Long Jump',
        'Lunges',
        'Military Parade',
        'Mixing',
        'Mopping Floor',
        'Nunchucks',
        'Parallel Bars',
        'Pizza Tossing',
        'Playing Cello',
        'Playing Daf',
        'Playing Dhol',
        'Playing Flute',
        'Playing Guitar',
        'Playing Piano',
        'Playing Sitar',
        'Playing Tabla',
        'Playing Violin',
        'Pole Vault',
        'Pommel Horse',
        'Pull Ups',
        'Punch',
        'Push Ups',
        'Rafting',
        'Rock Climbing Indoor',
        'Rope Climbing',
        'Rowing',
        'Salsa Spin',
        'Shaving Beard',
        'Shotput',
        'Skate Boarding',
        'Skiing',
        'Skijet',
        'Sky Diving',
        'Soccer Juggling',
        'Soccer Penalty',
        'Still Rings',
        'Sumo Wrestling',
        'Surfing',
        'Swing',
        'Table Tennis Shot',
        'Tai Chi',
        'Tennis Swing',
        'Throw Discus',
        'Trampoline Jumping',
        'Typing',
        'Uneven Bars',
        'Volleyball Spiking',
        'Walking With Dog',
        'Wall Pushups',
        'Writing On Board',
        'Yo Yo',
    ],
    'voc': [
        'aeroplane',
        'bicycle',
        'bird',
        'boat',
        'bottle',
        'bus',
        'car',
        'cat',
        'chair',
        'cow',
        'diningtable',
        'dog',
        'horse',
        'motorbike',
        'person',
        'pottedplant',
        'sheep',
        'sofa',
        'train',
        'tvmonitor',
    ],
}
_ZEROSHOT_CLASS_NAMES['imagenet_r'] = [
    _ZEROSHOT_CLASS_NAMES['imagenet'][i] for i in _IMAGENET_R_LABELSET
]
_ZEROSHOT_CLASS_NAMES['imagenet_a'] = [
    _ZEROSHOT_CLASS_NAMES['imagenet'][i] for i in _IMAGENET_A_LABELSET
]


def get_zeroshot_class_names(classnames_key):
  return _ZEROSHOT_CLASS_NAMES[classnames_key]


_ZEROSHOT_TEMPLATES = {
    'imagenet': [
        'a bad photo of a {}.',
        'a photo of many {}.',
        'a sculpture of a {}.',
        'a photo of the hard to see {}.',
        'a low resolution photo of the {}.',
        'a rendering of a {}.',
        'graffiti of a {}.',
        'a bad photo of the {}.',
        'a cropped photo of the {}.',
        'a tattoo of a {}.',
        'the embroidered {}.',
        'a photo of a hard to see {}.',
        'a bright photo of a {}.',
        'a photo of a clean {}.',
        'a photo of a dirty {}.',
        'a dark photo of the {}.',
        'a drawing of a {}.',
        'a photo of my {}.',
        'the plastic {}.',
        'a photo of the cool {}.',
        'a close-up photo of a {}.',
        'a black and white photo of the {}.',
        'a painting of the {}.',
        'a painting of a {}.',
        'a pixelated photo of the {}.',
        'a sculpture of the {}.',
        'a bright photo of the {}.',
        'a cropped photo of a {}.',
        'a plastic {}.',
        'a photo of the dirty {}.',
        'a jpeg corrupted photo of a {}.',
        'a blurry photo of the {}.',
        'a photo of the {}.',
        'a good photo of the {}.',
        'a rendering of the {}.',
        'a {} in a video game.',
        'a photo of one {}.',
        'a doodle of a {}.',
        'a close-up photo of the {}.',
        'a photo of a {}.',
        'the origami {}.',
        'the {} in a video game.',
        'a sketch of a {}.',
        'a doodle of the {}.',
        'a origami {}.',
        'a low resolution photo of a {}.',
        'the toy {}.',
        'a rendition of the {}.',
        'a photo of the clean {}.',
        'a photo of a large {}.',
        'a rendition of a {}.',
        'a photo of a nice {}.',
        'a photo of a weird {}.',
        'a blurry photo of a {}.',
        'a cartoon {}.',
        'art of a {}.',
        'a sketch of the {}.',
        'a embroidered {}.',
        'a pixelated photo of a {}.',
        'itap of the {}.',
        'a jpeg corrupted photo of the {}.',
        'a good photo of a {}.',
        'a plushie {}.',
        'a photo of the nice {}.',
        'a photo of the small {}.',
        'a photo of the weird {}.',
        'the cartoon {}.',
        'art of the {}.',
        'a drawing of the {}.',
        'a photo of the large {}.',
        'a black and white photo of a {}.',
        'the plushie {}.',
        'a dark photo of a {}.',
        'itap of a {}.',
        'graffiti of the {}.',
        'a toy {}.',
        'itap of my {}.',
        'a photo of a cool {}.',
        'a photo of a small {}.',
        'a tattoo of the {}.',
    ],
    # https://github.com/google-research/big_vision/blob/main/big_vision/evaluators/proj/image_text/prompt_engineering_constants.py.
    'imagenet_best': [
        'itap of a {}.',
        'a bad photo of the {}.',
        'a origami {}.',
        'a photo of the large {}.',
        'a {} in a video game.',
        'art of the {}.',
        'a photo of the small {}.',
        '{}',
    ],
    'birdsnap': ['a photo of a {}, a type of bird.'],
    'caltech101': [
        'a photo of a {}.',
        'a painting of a {}.',
        'a plastic {}.',
        'a sculpture of a {}.',
        'a sketch of a {}.',
        'a tattoo of a {}.',
        'a toy {}.',
        'a rendition of a {}.',
        'a embroidered {}.',
        'a cartoon {}.',
        'a {} in a video game.',
        'a plushie {}.',
        'a origami {}.',
        'art of a {}.',
        'graffiti of a {}.',
        'a drawing of a {}.',
        'a doodle of a {}.',
        'a photo of the {}.',
        'a painting of the {}.',
        'the plastic {}.',
        'a sculpture of the {}.',
        'a sketch of the {}.',
        'a tattoo of the {}.',
        'the toy {}.',
        'a rendition of the {}.',
        'the embroidered {}.',
        'the cartoon {}.',
        'the {} in a video game.',
        'the plushie {}.',
        'the origami {}.',
        'art of the {}.',
        'graffiti of the {}.',
        'a drawing of the {}.',
        'a doodle of the {}.',
    ],
    'cars196': [
        'a photo of a {}.',
        'a photo of the {}.',
        'a photo of my {}.',
        'i love my {}!',
        'a photo of my dirty {}.',
        'a photo of my clean {}.',
        'a photo of my new {}.',
        'a photo of my old {}.',
    ],
    'cifar10': [
        'a photo of a {}.',
        'a blurry photo of a {}.',
        'a black and white photo of a {}.',
        'a low contrast photo of a {}.',
        'a high contrast photo of a {}.',
        'a bad photo of a {}.',
        'a good photo of a {}.',
        'a photo of a small {}.',
        'a photo of a big {}.',
        'a photo of the {}.',
        'a blurry photo of the {}.',
        'a black and white photo of the {}.',
        'a low contrast photo of the {}.',
        'a high contrast photo of the {}.',
        'a bad photo of the {}.',
        'a good photo of the {}.',
        'a photo of the small {}.',
        'a photo of the big {}.',
    ],
    'cifar100': [
        'a photo of a {}.',
        'a blurry photo of a {}.',
        'a black and white photo of a {}.',
        'a low contrast photo of a {}.',
        'a high contrast photo of a {}.',
        'a bad photo of a {}.',
        'a good photo of a {}.',
        'a photo of a small {}.',
        'a photo of a big {}.',
        'a photo of the {}.',
        'a blurry photo of the {}.',
        'a black and white photo of the {}.',
        'a low contrast photo of the {}.',
        'a high contrast photo of the {}.',
        'a bad photo of the {}.',
        'a good photo of the {}.',
        'a photo of the small {}.',
        'a photo of the big {}.',
    ],
    'clevr_closest_lit_all': [
        '{} objects',
        '{} shapes',
        '{} thing',
        'the nearest shape in this image is {}.',
        'the closest shape in this rendered image is {}.',
        'the closest shape in this image is {}.',
    ],
    'clevr_counts': ['a photo of {} objects.'],
    'clevr_counts_lit_all': [
        '{} objects',
        '{} things',
        'a photo of {} objects.',
        'a picture of {} objects.',
        'there are {} objects.',
        'there are {} objects in the image.',
        'a rendered image of {} objects.',
        'a picture of {} shapes.',
        'there are {} shapes.',
        'there are {} shapes in the image.',
        'a rendered image of {} shapes.',
    ],
    'country211': [
        'a photo i took in {}.',
        'a photo i took while visiting {}.',
        'a photo from my home country of {}.',
        'a photo from my visit to {}.',
        'a photo showing the country of {}.',
    ],
    'dsprites_pos_lit_all': [
        'an object located {}'
    ],
    'dsprites_orient_lit_all': [
        'an object rotated at {}',
        'something rotated at {}',
        '{} rotation',
        'something at a {} rotation',
    ],
    'dtd': [
        'a photo of a {} texture.',
        'a photo of a {} pattern.',
        'a photo of a {} thing.',
        'a photo of a {} object.',
        'a photo of the {} texture.',
        'a photo of the {} pattern.',
        'a photo of the {} thing.',
        'a photo of the {} object.',
    ],
    'dtd_lit_all': ['a {} texture.'],
    'eurosat': [
        'a centered satellite photo of {}.',
        'a centered satellite photo of a {}.',
        'a centered satellite photo of the {}.',
    ],
    'eurosat_lit_all': [
        'a satellite photo of {}.',
        '{} from above.',
        'an aerial view of {}.',
        'an overhead view of {}.',
        'an overhead image of {}.',
        'a satellite image of {}.',
        'photo of {} from the sky.',
    ],
    'facial_emotion_recognition_2013': [
        'a photo of a {} looking face.',
        'a photo of a face showing the emotion: {}.',
        'a photo of a face looking {}.',
        'a face that looks {}.',
        'they look {}.',
        'look at how {} they are.',
    ],
    'fgvc_aircraft': [
        'a photo of a {}, a type of aircraft.',
        'a photo of the {}, a type of aircraft.',
    ],
    'food101': ['a photo of {}, a type of food.'],
    'gtsrb': [
        'a zoomed in photo of a "{}" traffic sign.',
        'a centered photo of a "{}" traffic sign.',
        'a close up photo of a "{}" traffic sign.',
    ],
    'hateful_memes': ['a {}.'],
    'kitti': ['{}'],
    'kinetics700': [
        'a photo of {}.',
        'a photo of a person {}.',
        'a photo of a person using {}.',
        'a photo of a person doing {}.',
        'a photo of a person during {}.',
        'a photo of a person performing {}.',
        'a photo of a person practicing {}.',
        'a video of {}.',
        'a video of a person {}.',
        'a video of a person using {}.',
        'a video of a person doing {}.',
        'a video of a person during {}.',
        'a video of a person performing {}.',
        'a video of a person practicing {}.',
        'a example of {}.',
        'a example of a person {}.',
        'a example of a person using {}.',
        'a example of a person doing {}.',
        'a example of a person during {}.',
        'a example of a person performing {}.',
        'a example of a person practicing {}.',
        'a demonstration of {}.',
        'a demonstration of a person {}.',
        'a demonstration of a person using {}.',
        'a demonstration of a person doing {}.',
        'a demonstration of a person during {}.',
        'a demonstration of a person performing {}.',
        'a demonstration of a person practicing {}.',
    ],
    'mnist': ['a photo of the number: "{}".'],
    'none': ['{}'],
    'oxford_flowers102': ['a photo of a {}, a type of flower.'],
    'oxford_flowers102_lit_all': ['a {} flower', 'a {} plant'],
    'oxford_iiit_pet': ['a photo of a {}, a type of pet.'],
    'oxford_iiit_pet_lit_all': [
        'a type of pet {}', 'a {} texture', '{}, an animal'
    ],
    'patch_camelyon': ['this is a photo of {}'],
    'patch_camelyon_lit_all': [
        'a histopathology slide showing {}',
        'histopathology image of {}',
        'an example of {}',
        'an example histopathological image showing {}',
        "patient's pathology examination indicates {}",
        'a {} slide',
    ],
    'photo': ['a photo of a {}.'],
    'resisc45': [
        'satellite imagery of {}.',
        'aerial imagery of {}.',
        'satellite photo of {}.',
        'aerial photo of {}.',
        'satellite view of {}.',
        'aerial view of {}.',
        'satellite imagery of a {}.',
        'aerial imagery of a {}.',
        'satellite photo of a {}.',
        'aerial photo of a {}.',
        'satellite view of a {}.',
        'aerial view of a {}.',
        'satellite imagery of the {}.',
        'aerial imagery of the {}.',
        'satellite photo of the {}.',
        'aerial photo of the {}.',
        'satellite view of the {}.',
        'aerial view of the {}.',
    ],
    'resisc45_lit_all': [
        'a satellite photo of {}.',
        '{} from above.',
        'an aerial view of {}.',
        'a satellite image of {}.',
    ],
    'retino_lit_all': [
        'a retinal image with {}',
        'a retina with {}',
        'a fundus image with signs of {}',
    ],
    'smallnorb_azmth_lit_all': [
        'an object rotated at {}',
        'something rotated at {}',
        '{} rotation',
        'something at a {} rotation',
    ],
    'smallnorb_elev_lit_all': [
        'an object rotated at {}',
        'something rotated at {}',
        '{} rotation',
        'something at a {} rotation',
    ],
    'sst2': ['a {} review of a movie.'],
    'stl10': [
        'a photo of a {}.',
        'a photo of the {}.',
    ],
    'sun397': [
        'a photo of a {}.',
        'a photo of the {}.',
    ],
    'svhn_lit_all': [
        'the number {}',
        'a street sign with the number {}',
        'a photo of the number {} written on a sign',
        'an outdoor house number {}',
        'the number {} in the center of the image',
        'an outdoor number {} written on a sign',
        'a centered image of the number {}'
    ],
    'ucf101': [
        'a photo of a person {}.',
        'a video of a person {}.',
        'a example of a person {}.',
        'a demonstration of a person {}.',
        'a photo of the person {}.',
        'a video of the person {}.',
        'a example of the person {}.',
        'a demonstration of the person {}.',
        'a photo of a person using {}.',
        'a video of a person using {}.',
        'a example of a person using {}.',
        'a demonstration of a person using {}.',
        'a photo of the person using {}.',
        'a video of the person using {}.',
        'a example of the person using {}.',
        'a demonstration of the person using {}.',
        'a photo of a person doing {}.',
        'a video of a person doing {}.',
        'a example of a person doing {}.',
        'a demonstration of a person doing {}.',
        'a photo of the person doing {}.',
        'a video of the person doing {}.',
        'a example of the person doing {}.',
        'a demonstration of the person doing {}.',
        'a photo of a person during {}.',
        'a video of a person during {}.',
        'a example of a person during {}.',
        'a demonstration of a person during {}.',
        'a photo of the person during {}.',
        'a video of the person during {}.',
        'a example of the person during {}.',
        'a demonstration of the person during {}.',
        'a photo of a person performing {}.',
        'a video of a person performing {}.',
        'a example of a person performing {}.',
        'a demonstration of a person performing {}.',
        'a photo of the person performing {}.',
        'a video of the person performing {}.',
        'a example of the person performing {}.',
        'a demonstration of the person performing {}.',
        'a photo of a person practicing {}.',
        'a video of a person practicing {}.',
        'a example of a person practicing {}.',
        'a demonstration of a person practicing {}.',
        'a photo of the person practicing {}.',
        'a video of the person practicing {}.',
        'a example of the person practicing {}.',
        'a demonstration of the person practicing {}.',
    ],
    'voc': ['a photo of a {}.'],
}


def get_zeroshot_template(prompts_key):
  return _ZEROSHOT_TEMPLATES[prompts_key]
